
import time
#RabbitMQ
import pika
#多线程
from threading import Thread

#自定义的RabbitMQ类
class RabbitUtil:
    #初始化
    def __init__(self, host, user, passwd, port=5672, vhost='/'):
        #相关参数存储
        self.host    = host;
        self.user    = user;
        self.passwd  = passwd;
        self.port    = port;
        self.vhost   = vhost;
        #---------内部使用---------
        self.conn    = None;
        self.channel = None;
        #初始化连接
        self.Connect();
        #end
    
    #初始化Rabbit的连接
    def Connect(self):
        #建立连接
        UserInfo  = pika.PlainCredentials(self.user, self.passwd);
        self.conn = pika.BlockingConnection(pika.ConnectionParameters(self.host, self.port, self.vhost, UserInfo, heartbeat=600));
        #end

    #关闭Rabbit的连接
    def Close(self):
        if self.conn:
            self.conn.close();
            self.conn = None;
        #end

    #保持连接
    def ReConnect(self):
        if self.conn == None or self.conn.is_closed:
            self.Connect();
        #end

    #重新开启Channel(包括了ReConnect)
    def ReChannel(self):
        #重新建立连接
        self.ReConnect();
        #重新创建Channel
        if self.channel == None:
            self.channel = self.conn.channel();
        #end

    #创建队列
    def CreateQueues(self, QueueList):
        #建立Channel
        self.ReChannel();
        #通过Channel，创建Queue
        for OneQueue in QueueList:
            # 如果指定的queue不存在，则会创建一个queue，如果已经存在 则不会做其他动作，
            # 官方推荐，每次使用时都可以加上这句
            self.channel.queue_declare(queue=OneQueue, durable=True); #durable持久化数据
            # print("Create Queue:", OneQueue);
        #end for
        #end

    #获得指定队列的当前长度
    def QueueLen(self, QueueName):
        #建立Channel
        self.ReChannel();
        #
        res = self.channel.queue_declare(queue=QueueName, durable=True);
        return res.method.message_count;
        #end

    def queuePurge(self, QueueName):
        #建立Channel
        self.ReChannel();
        #
        self.channel.queue_purge(queue=QueueName);
        
        
    # 发送消息列表
    def Produce(self, QueueName, MsgList):
        #建立Channel
        self.ReChannel();
        #依次发送每条消息
        for Msg in MsgList:
            # print("Insert into", QueueName, "=", Msg);
            #通过通道发送消息
            self.channel.basic_publish(exchange = '',           # 当前是一个简单模式，所以这里设置为空字符串就可以了
                                       routing_key = QueueName, # 指定消息要发送到哪个queue
                                       body = Msg,              # 指定要发送的消息
                                       properties = pika.BasicProperties(delivery_mode=2)); #消息持久化
        #end for
        #end

    #注册消费者处理函数 并阻塞
    def RegistConsumer(self, QueueName, ConsumerFunc):
        # 回调函数
        def callback(ch, method, properties, body):
            #print('Consume Msg:', body);
            #消息到来时执行消费函数
            Ret = ConsumerFunc(body);
            #成功处理消息，才将消息从队列中删除
            if Ret == 0:
                #消费后，从队列中删除
                ch.basic_ack(delivery_tag=method.delivery_tag);
            #end

        #确保Channel
        self.ReChannel();
        # prefetch_count表示接收的消息数量，当我接收的消息没有处理完（用basic_ack标记消息已处理完毕）之前不会再接收新的消息了
        # 还有就是这个设置必须在basic_consume之上，否则不生效
        self.channel.basic_qos(prefetch_count=1);
        # channel:    包含channel的一切属性和方法
        # method:     包含 consumer_tag, delivery_tag, exchange, redelivered, routing_key
        # properties: basic_publish 通过 properties 传入的参数
        # body:       basic_publish发送的消息
        self.channel.basic_consume(queue=QueueName,                 # 接收指定queue的消息
                                   auto_ack=False,                  # 指定为False，处理完消息后，回调中删除消息
                                   on_message_callback=callback);   # 设置收到消息的回调函数

        # 一直处于等待接收消息的状态，如果没收到消息就一直处于阻塞状态，收到消息就调用上面的回调函数
        self.channel.start_consuming();
        #end


#生产者任务 for test
def ProduceTask():
    #初始化RabbitMQ操作对象
    RMQ = RabbitUtil(host="127.0.0.1", user="user", passwd="123456", vhost='/');
    #创建Queue
    RMQ.CreateQueues(["QueueTest"]);
    #模拟消息发送
    MsgIdx = 0;
    while True:
        #发送消息
        RMQ.Produce("QueueTest", [f"http://www.163.com/{MsgIdx}"]);
        #消息序号自增加
        MsgIdx = MsgIdx+1;
        #每3S发送一条消息
        time.sleep(3);
    #end while

#消费任务 for test
def ConsumeTask():
    #测试用消费函数
    def ConsumerFunc(Msg):
        print("Consume Msg :", Msg);
        #默认延时处理
        time.sleep(2);
        #输出结果
        return 0;
        #end
    
    #初始化RabbitMQ操作对象
    RMQ = RabbitUtil(host="127.0.0.1", user="user", passwd="123456", vhost='/');
    #注册消费者
    RMQ.RegistConsumer("QueueTest", ConsumerFunc);
    #Can't reach here
    return 0;
    #end


#测试

if __name__ == "__main__":
    ThreadList = [];
    #启动一个生产者
    TH = Thread(target=ProduceTask, args=());
    TH.start();
    ThreadList.append(TH);
    
    #延时启动两个消费者
    time.sleep(50);
    #消费者1
    TH = Thread(target=ConsumeTask, args=());
    TH.start();
    ThreadList.append(TH);
    #消费者2
    TH = Thread(target=ConsumeTask, args=());
    TH.start();
    ThreadList.append(TH);

    #等待每个线程的退出
    for TH in ThreadList:
        TH.join();
    #end for
    #end


